'''
返回不同类型的向量存储
'''
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore

from model_config import embedding_model_define

# MongoDB
class MongoDBConfig:
  def __init__(self):
    self.uri='mongodb+srv://beerspume:azXjrMzzcq9EUDCA@cluster-rag.k2ffs.mongodb.net/'
    self.dbname='AI'
    self.collection='RAG'
    self.index_name='langchain-index-vectorstores'
mongodb_config:MongoDBConfig=MongoDBConfig()

def get_vector_storage_mongodb(embedding:Embeddings)->VectorStore:
  from langchain_mongodb import MongoDBAtlasVectorSearch
  from pymongo import MongoClient
  mongo_client=MongoClient(mongodb_config.uri)
  mongo_db_name=mongodb_config.dbname
  mongo_collection_name=mongodb_config.collection
  mongo_collection=mongo_client[mongo_db_name][mongo_collection_name]
  vector_store=MongoDBAtlasVectorSearch(
    collection=mongo_collection,
    embedding=embedding,
    index_name=mongodb_config.index_name,
    relevance_score_fn='cosine'
  )
  try:
    vector_store.create_vector_search_index(
      dimensions=embedding_model_define.dimensions,
    )
  except:
    pass
  return vector_store

# Chroma
class ChromaDBConfig:
  def __init__(self):
    self.path='./data/test_chroma'
    self.collection='RAG'
    self.index_name='langchain-index-vectorstores'
chromadb_config:ChromaDBConfig=ChromaDBConfig()

def get_vector_storage_chroma(embedding:Embeddings)->VectorStore:
  from langchain_chroma.vectorstores import Chroma
  import chromadb
  chroma_client=chromadb.PersistentClient(path=chromadb_config.path)
  vector_store=Chroma(
    chromadb_config.collection,
    embedding_function=embedding,
    client=chroma_client,
  )
  return vector_store


def get_vector_storage_default(embedding:Embeddings)->VectorStore:
  return get_vector_storage_chroma(embedding=embedding)

